/*
 * @param M                 Faiss(IVFPQFS)搜索参数：子量化器数量
 * @param k_f               Faiss(IVFPQFS)构建参数：邻居节点数
 * @param nbits             Faiss(IVFPQFS)搜索参数：每个子量化器的位数（固定为4）
 * @param nprobe            Faiss(IVFPQFS)构建参数：倒排列表（Inverted List）数量
 * @param bbs               Faiss(IVFPQFS)构建参数：预计算表块大小（32的倍数）
 * @param qbs               Faiss(IVFPQFS)构建参数：向量量化位数（KPRetrieval可选）
 * @param metric            Faiss(IVFPQFS)构建参数：距离度量类型
 * @param quantizer         Faiss(IVFPQFS)量化器
 * @param index             Faiss(IVFPQFS)量化索引
 * @param refine_index      Faiss(IVFPQFS)精准索引
 * @param ivf_params        Faiss(IVFPQFS)ivf参数：用于添加nprobe
 * @param refine_params     Faiss(IVFPQFS)refine参数：用于添加k_f
 */

#ifndef FAISS_IVFPQFSFS_INDEX_H
#define FAISS_IVFPQFSFS_INDEX_H

#include "index_base.h"
#include "datasets.h"
#include "config_parser.h"
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFPQFastScan.h>
#include <faiss/IndexRefine.h>
#include <memory>

class FaissIVFPQFSIndex : public IndexBase {
public:
    FaissIVFPQFSIndex(const ConfigParser& config);
    ~FaissIVFPQFSIndex();
    virtual void build(const Datasets* data) override;
    virtual void prepare_search() override;
    virtual void search_single(const float* xq, int k, float* distances, int64_t* labels) override;
    virtual void search_batch(const int nq, const float* xq, int k, float* distances, int64_t* labels) override;

private:
    int M;
    int k_f;
    int nbits;
    int nprobe;
    int bbs;
#ifdef QUANT
    int qbs;
#endif
    faiss::MetricType metric;
    std::unique_ptr<faiss::IndexFlat> quantizer;
    std::unique_ptr<faiss::IndexIVFPQFastScan> index;
    std::unique_ptr<faiss::IndexRefineFlat> refine_index;
    faiss::SearchParametersIVF ivf_params;
    faiss::IndexRefineSearchParameters refine_params;
    std::string index_path;
    std::string save_or_load;
};

#endif // FAISS_IVFPQFS_INDEX_H